from app.extensions import db
from datetime import datetime
from marshmallow import Schema, fields

class Relationship(db.Model):
    __tablename__ = 'relationships'

    id = db.Column(db.Integer, primary_key=True)
    chart_id = db.Column(db.Integer, db.ForeignKey('charts.id'), nullable=False)
    upstream_node_id = db.Column(db.Integer, db.ForeignKey('nodes.id'), nullable=False)
    downstream_node_id = db.Column(db.Integer, db.ForeignKey('nodes.id'), nullable=False)
    relationship_type = db.Column(db.String(50), default='supply')
    label = db.Column(db.String(100))
    strength = db.Column(db.Numeric(3, 2), default=1.0)
    color = db.Column(db.String(20), default='#67C23A')
    status = db.Column(db.Enum('active', 'inactive'), default='active')
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    __table_args__ = (
        db.UniqueConstraint('chart_id', 'upstream_node_id', 'downstream_node_id', 'relationship_type'),
    )

    def to_dict(self):
        return {
            'id': self.id,
            'chart_id': self.chart_id,
            'upstream_node_id': self.upstream_node_id,
            'downstream_node_id': self.downstream_node_id,
            'relationship_type': self.relationship_type,
            'label': self.label,
            'strength': float(self.strength) if self.strength else 1.0,
            'color': self.color,
            'status': self.status
        }

class RelationshipSchema(Schema):
    id = fields.Int(dump_only=True)
    chart_id = fields.Int(required=True)
    upstream_node_id = fields.Int(required=True)
    downstream_node_id = fields.Int(required=True)
    relationship_type = fields.Str()
    label = fields.Str()
    strength = fields.Float()
    color = fields.Str()
    status = fields.Str()